{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9cee9a09",
   "metadata": {},
   "source": [
    "# 注解 Target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3ef0ac3f",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "000acf47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "\n",
    "import tvm\n",
    "import tvm.relay.testing\n",
    "import tvm.relay.transform as transform\n",
    "from tvm import relay\n",
    "from tvm import runtime\n",
    "from tvm.contrib import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f46caac",
   "metadata": {},
   "source": [
    "## 注解多 end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64cc4f8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.test\")\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    # 注册 nn.relu 算子在 test 目标上的处理函数\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "18a5f567",
   "metadata": {},
   "outputs": [],
   "source": [
    "def before():\n",
    "    x = relay.var(\"x\", shape=(10, 10))\n",
    "    r = relay.nn.relu(x)\n",
    "    a_1 = relay.abs(r)\n",
    "    a_2 = relay.abs(r)\n",
    "    out = relay.add(a_1, a_2)\n",
    "    f = relay.Function([x], out)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    return mod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd5aefd3",
   "metadata": {},
   "source": [
    "`relu` 在 `test` 目标上执行，其余算子在 `default` 目标上执行："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "41e61cd9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span>x, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> abs(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #A2F; font-weight: bold\">=</span> abs(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">12</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 测试两种场景：不注解非调用操作和注解非调用操作\n",
    "for annotate_non_call_ops in [False, True]:\n",
    "    # 应用 AnnotateTarget 转换\n",
    "    result = transform.AnnotateTarget(\"test\", annotate_non_call_ops)(before())\n",
    "    if annotate_non_call_ops:\n",
    "        result.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59bcacfa",
   "metadata": {},
   "source": [
    "## 类型传播"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b57e1dc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "target = \"test_type_propagation\"\n",
    "\n",
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.\" + target)\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return expr.args[0].checked_type.dtype == \"float32\"\n",
    "\n",
    "def before():\n",
    "    x = relay.var(\"x\", shape=(10, 10))\n",
    "    r = relay.nn.relu(x)\n",
    "    out = relay.nn.relu(r)\n",
    "    f = relay.Function([x], out)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    return mod\n",
    "\n",
    "for annotate_non_call_ops in [False, True]:\n",
    "    # If the type isn't propogated, then the relu checker function will fail to get the dtype.\n",
    "    assert transform.AnnotateTarget(target, annotate_non_call_ops)(before())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0f4d0925",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span>x, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_type_propagation&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_type_propagation&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_type_propagation&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_type_propagation&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "transform.AnnotateTarget(target, True)(before()).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "415109ef",
   "metadata": {},
   "source": [
    "## RefCreate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6fe8b69e",
   "metadata": {},
   "outputs": [],
   "source": [
    "target = \"relu\"\n",
    "\n",
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.\" + target)\n",
    "def annotate(expr):\n",
    "    return True\n",
    "\n",
    "def before():\n",
    "    ref = relay.expr.RefCreate(relay.const(1.0))\n",
    "    r = relay.expr.RefWrite(ref, relay.nn.relu(relay.expr.RefRead(ref)))\n",
    "    return tvm.IRModule.from_expr(r)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c8a2c42a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>() <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> () {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #008000\">1</span>f <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> ref(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>);\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>ref(float32) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>ref(float32) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>ref(float32) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>ref(float32) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #A2F; font-weight: bold\">=</span> ref_read(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>);\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;relu&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;relu&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">12</span> <span style=\"color: #A2F; font-weight: bold\">=</span> ref_write(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>);\n",
       "  annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>() <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "result.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6e35fda",
   "metadata": {},
   "source": [
    "## 测试 TupleNode 在 AnnotateTarget 转换中的处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c836c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义目标编译器标识\n",
    "target = \"test_tuple\"\n",
    "\n",
    "# 注册 nn.relu 操作在指定目标上的处理函数\n",
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.\" + target)\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "# 注册 concatenate 操作在指定目标上的处理函数\n",
    "@tvm.ir.register_op_attr(\"concatenate\", \"target.\" + target)\n",
    "def concatenate(expr):  # pylint: disable=unused-variable\n",
    "    return True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa06d290",
   "metadata": {},
   "source": [
    " 验证当元组被支持的节点包围时，元组是否被包含在注解"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dfc4ff15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span>x, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span>y, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #A2F; font-weight: bold\">=</span> (<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">5</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #A2F; font-weight: bold\">=</span> concatenate(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, axis<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test_tuple&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    x = relay.var(\"x\", shape=(10, 5))\n",
    "    y = relay.var(\"y\", shape=(10, 5))\n",
    "    a_1 = relay.nn.relu(x)\n",
    "    a_2 = relay.nn.relu(y)\n",
    "    out = relay.concatenate((a_1, a_2), axis=1)\n",
    "    f = relay.Function([x, y], out)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    return mod\n",
    "result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "result.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd854664",
   "metadata": {},
   "source": [
    "## 测试复合函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c0127a9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>a: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>b: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span>a, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #A2F; font-weight: bold\">=</span> annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_begin(<span style=\"color: #A2F; font-weight: bold\">%</span>b, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>in_1: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>in_2: Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>, Composite<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;test.add_relu&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] {\n",
       "    <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> add(<span style=\"color: #A2F; font-weight: bold\">%</span>in_1, <span style=\"color: #A2F; font-weight: bold\">%</span>in_2) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "    nn<span style=\"color: #A2F; font-weight: bold\">.</span>relu(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32], Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32]) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  annotation<span style=\"color: #A2F; font-weight: bold\">.</span>compiler_end(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>, compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;default&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">10</span>, <span style=\"color: #008000\">10</span>), float32] <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def before():\n",
    "    a = relay.var(\"a\", shape=(10, 10))\n",
    "    b = relay.var(\"b\", shape=(10, 10))\n",
    "\n",
    "    # add_relu function\n",
    "    in_1 = relay.var(\"in_1\", shape=(10, 10))\n",
    "    in_2 = relay.var(\"in_2\", shape=(10, 10))\n",
    "    add_node = relay.add(in_1, in_2)\n",
    "    relu_node = relay.nn.relu(add_node)\n",
    "    add_relu = relay.Function([in_1, in_2], relu_node)\n",
    "    add_relu = add_relu.with_attr(\"Composite\", \"test.add_relu\")\n",
    "\n",
    "    # merged function\n",
    "    r = relay.Call(add_relu, [a, b])\n",
    "    f = relay.Function([a, b], r)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    return mod\n",
    "result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "result.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f00dc94",
   "metadata": {},
   "source": [
    "## 测试双目标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e4831222",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.double.A\")\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def before():\n",
    "    x = relay.var(\"x\", shape=(10, 5))\n",
    "    a_1 = relay.nn.relu(x)\n",
    "    mod = tvm.IRModule.from_expr(a_1)\n",
    "    return mod\n",
    "\n",
    "for annotate_non_call_ops in [True, False]:\n",
    "    mod = before()\n",
    "    mod1 = transform.AnnotateTarget(\"double.A\", annotate_non_call_ops)(mod)\n",
    "    mod2 = transform.AnnotateTarget(\"double.A\", annotate_non_call_ops)(mod1)\n",
    "    tvm.ir.assert_structural_equal(mod1, mod2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b0c257f",
   "metadata": {},
   "source": [
    "## 测试不同目标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "da97838c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.different.A\")\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "@tvm.ir.register_op_attr(\"add\", \"target.different.B\")\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def before():\n",
    "    x = relay.var(\"x\", shape=(10, 5))\n",
    "    a_1 = relay.nn.relu(x)\n",
    "    b_1 = relay.add(a_1, a_1)\n",
    "    mod = tvm.IRModule.from_expr(b_1)\n",
    "    return mod\n",
    "\n",
    "for annotate_non_call_ops in [True, False]:\n",
    "    mod = before()\n",
    "    mod1 = transform.AnnotateTarget(\"different.A\", annotate_non_call_ops)(mod)\n",
    "    mod1 = transform.AnnotateTarget(\"different.B\", annotate_non_call_ops)(mod1)\n",
    "    mod2 = transform.AnnotateTarget([\"different.A\", \"different.B\"], annotate_non_call_ops)(mod)\n",
    "    tvm.ir.assert_structural_equal(mod1, mod2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ccf268c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.ir.register_op_attr(\"nn.relu\", \"target.A\")\n",
    "def relu(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "@tvm.ir.register_op_attr(\"add\", \"target.B\")\n",
    "def add(expr):  # pylint: disable=unused-variable\n",
    "    return True\n",
    "\n",
    "def before():\n",
    "    x = relay.var(\"x\", shape=(10, 5))\n",
    "    a_1 = relay.nn.relu(x)\n",
    "    a_2 = relay.abs(a_1)\n",
    "    a_3 = relay.nn.relu(a_1)\n",
    "    out = relay.add(a_2, a_3)\n",
    "\n",
    "    f = relay.Function([x], out)\n",
    "    mod = tvm.IRModule.from_expr(f)\n",
    "    return mod\n",
    "\n",
    "for annotate_non_call_ops in [True, False]:\n",
    "    mod = transform.AnnotateTarget(\"A\", annotate_non_call_ops)(before())\n",
    "    mod = transform.AnnotateTarget(\"B\", annotate_non_call_ops)(mod)\n",
    "    expected = transform.AnnotateTarget([\"A\", \"B\"], annotate_non_call_ops)(before())\n",
    "    tvm.ir.assert_structural_equal(expected, mod)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "891153ed",
   "metadata": {},
   "source": [
    "## 其他"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5a8fb60e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_ends_with_tuple():\n",
    "    \"\"\"测试以元组结尾的计算图在AnnotateTarget转换中的处理\n",
    "\n",
    "    此测试用例验证当计算图以元组（Tuple）结尾，并且包含TupleGetItem操作时，\n",
    "    AnnotateTarget转换能否正确地处理这些结构并添加适当的编译器注解。\n",
    "    \"\"\"\n",
    "    # 定义目标编译器标识\n",
    "    trgt = \"clip\"\n",
    "\n",
    "    # 注册clip操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"clip\", \"target.\" + trgt)\n",
    "    def relu(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    def get_model(get_item):\n",
    "        \"\"\"创建测试模型\n",
    "\n",
    "        参数:\n",
    "            get_item: 布尔值，表示是否使用TupleGetItem操作获取元组中的元素\n",
    "\n",
    "        返回:\n",
    "            包含clip操作和元组结构的IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        a = relay.var(\"a\", shape=(1, 16, 16, 4), dtype=\"uint8\")\n",
    "        # 应用clip操作\n",
    "        z = relay.op.clip(a, 0, 255)\n",
    "        # 对同一输入应用不同范围的clip操作\n",
    "        b = relay.op.clip(z, 0, 15)\n",
    "        c = relay.op.clip(z, 16, 31)\n",
    "        # 创建包含两个clip结果的元组\n",
    "        t = relay.Tuple((c, b))\n",
    "        # 根据参数决定是否使用TupleGetItem获取元组中的第二个元素\n",
    "        tgi = relay.TupleGetItem(t, 1) if get_item else t\n",
    "        foo = relay.Function([a], tgi)\n",
    "        # 返回创建的IR模块\n",
    "        return tvm.IRModule.from_expr(tgi)\n",
    "\n",
    "    def get_expected(annotate_non_call_ops, get_item):\n",
    "        \"\"\"生成预期的注解后模型\n",
    "\n",
    "        参数:\n",
    "            annotate_non_call_ops: 布尔值，表示是否为非调用操作添加注解\n",
    "            get_item: 布尔值，表示是否使用TupleGetItem操作\n",
    "\n",
    "        返回:\n",
    "            预期的包含正确编译器注解的IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        a_ = relay.var(\"a\", shape=(1, 16, 16, 4), dtype=\"uint8\")\n",
    "        # 为输入添加编译器开始注解\n",
    "        a = relay.annotation.compiler_begin(a_, trgt)\n",
    "        # 应用clip操作\n",
    "        z = relay.op.clip(a, 0, 255)\n",
    "        # 为第一个clip结果添加编译器结束注解\n",
    "        z1 = relay.annotation.compiler_end(z, trgt)\n",
    "        # 为z1添加编译器开始注解\n",
    "        z1 = relay.annotation.compiler_begin(z1, trgt)\n",
    "        # 应用第二个clip操作\n",
    "        b = relay.op.clip(z1, 0, 15)\n",
    "        # 为b添加编译器结束注解\n",
    "        b = relay.annotation.compiler_end(b, trgt)\n",
    "        # 根据参数决定是否为b添加编译器开始注解\n",
    "        b = relay.annotation.compiler_begin(b, trgt) if annotate_non_call_ops else b\n",
    "        # 为z添加编译器结束注解\n",
    "        z2 = relay.annotation.compiler_end(z, trgt)\n",
    "        # 为z2添加编译器开始注解\n",
    "        z2 = relay.annotation.compiler_begin(z2, trgt)\n",
    "        # 应用第三个clip操作\n",
    "        c = relay.op.clip(z2, 16, 31)\n",
    "        # 为c添加编译器结束注解\n",
    "        c = relay.annotation.compiler_end(c, trgt)\n",
    "        # 根据参数决定是否为c添加编译器开始注解\n",
    "        c = relay.annotation.compiler_begin(c, trgt) if annotate_non_call_ops else c\n",
    "        # 创建包含c和b的元组\n",
    "        t = relay.Tuple((c, b))\n",
    "        # 根据参数决定是否为元组添加编译器结束注解\n",
    "        t = relay.annotation.compiler_end(t, trgt) if annotate_non_call_ops else t\n",
    "        if get_item:\n",
    "            # 如果使用TupleGetItem，根据参数决定是否为元组添加编译器开始注解\n",
    "            t = relay.annotation.compiler_begin(t, trgt) if annotate_non_call_ops else t\n",
    "            # 获取元组中的第二个元素\n",
    "            tgi = relay.TupleGetItem(t, 1)\n",
    "            # 根据参数决定是否为TupleGetItem结果添加编译器结束注解\n",
    "            tgi = relay.annotation.compiler_end(tgi, trgt) if annotate_non_call_ops else tgi\n",
    "        else:\n",
    "            tgi = t\n",
    "        foo = relay.Function([a_], tgi)\n",
    "        # 返回预期的IR模块\n",
    "        return tvm.IRModule.from_expr(foo)\n",
    "\n",
    "    # 测试所有组合场景\n",
    "    for get_item in [True, False]:\n",
    "        for annotate_non_call_ops in [False, True]:\n",
    "            # 获取测试模型\n",
    "            mod = get_model(get_item)\n",
    "            # 应用AnnotateTarget转换\n",
    "            mod = transform.AnnotateTarget(\"clip\", annotate_non_call_ops)(mod)\n",
    "            # 获取预期结果并应用类型推断\n",
    "            expected = transform.InferType()(get_expected(annotate_non_call_ops, get_item))\n",
    "            # 验证实际结果与预期结果是否结构相等\n",
    "            tvm.ir.assert_structural_equal(expected, mod)\n",
    "\n",
    "\n",
    "def test_if_else():\n",
    "    \"\"\"测试条件分支（If-else）节点在AnnotateTarget转换中的处理\n",
    "\n",
    "    此测试用例验证当计算图包含条件分支结构，并且被支持的操作节点包围时，\n",
    "    AnnotateTarget转换能否正确地为这些结构添加编译器注解。\n",
    "    \"\"\"\n",
    "    # 定义目标编译器标识\n",
    "    target = \"test_if_else\"\n",
    "\n",
    "    # 注册equal操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"equal\", \"target.\" + target)\n",
    "    def relu(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册tanh操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"tanh\", \"target.\" + target)\n",
    "    def tanh(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册sigmoid操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"sigmoid\", \"target.\" + target)\n",
    "    def sigmoid(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册erf操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"erf\", \"target.\" + target)\n",
    "    def erf(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    def before():\n",
    "        \"\"\"创建包含条件分支的原始计算图\n",
    "\n",
    "        返回:\n",
    "            包含If-else结构的IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        data = relay.var(\"data\", shape=(1, 32))\n",
    "        eq1 = relay.var(\"e1\", shape=[], dtype=\"float32\")\n",
    "        eq2 = relay.var(\"e2\", shape=[], dtype=\"float32\")\n",
    "        # 创建条件表达式\n",
    "        eq = relay.equal(eq1, eq2)\n",
    "\n",
    "        # 创建条件为真时的分支（应用tanh激活函数）\n",
    "        true_branch = relay.tanh(data)\n",
    "        # 创建条件为假时的分支（应用sigmoid激活函数）\n",
    "        false_branch = relay.sigmoid(data)\n",
    "        # 创建条件分支结构\n",
    "        ife = relay.If(eq, true_branch, false_branch)\n",
    "        # 对条件分支的结果应用erf操作\n",
    "        out = relay.erf(ife)\n",
    "        func = relay.Function([data, eq1, eq2], out)\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "\n",
    "        return mod\n",
    "\n",
    "    def after():\n",
    "        \"\"\"生成预期的注解后计算图\n",
    "\n",
    "        返回:\n",
    "            包含正确编译器注解的条件分支IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        data = relay.var(\"data\", shape=(1, 32))\n",
    "        eq1 = relay.var(\"e1\", shape=[], dtype=\"float32\")\n",
    "        eq2 = relay.var(\"e2\", shape=[], dtype=\"float32\")\n",
    "\n",
    "        # 为条件变量添加编译器开始注解\n",
    "        cb_1 = relay.annotation.compiler_begin(eq1, target)\n",
    "        cb_2 = relay.annotation.compiler_begin(eq2, target)\n",
    "\n",
    "        # 创建条件表达式并添加编译器结束注解\n",
    "        equality_condition = relay.equal(cb_1, cb_2)\n",
    "        ce_1 = relay.annotation.compiler_end(equality_condition, target)\n",
    "\n",
    "        # 为条件为真的分支添加编译器注解\n",
    "        cb_3 = relay.annotation.compiler_begin(data, target)\n",
    "        true_branch = relay.tanh(cb_3)\n",
    "        ce_2 = relay.annotation.compiler_end(true_branch, target)\n",
    "\n",
    "        # 为条件为假的分支添加编译器注解\n",
    "        cb_4 = relay.annotation.compiler_begin(data, target)\n",
    "        false_branch = relay.sigmoid(cb_4)\n",
    "        ce_3 = relay.annotation.compiler_end(false_branch, target)\n",
    "\n",
    "        # 创建条件分支结构\n",
    "        if_condition = relay.If(ce_1, ce_2, ce_3)\n",
    "        # 为条件分支结果和最终erf操作添加编译器注解\n",
    "        cb_5 = relay.annotation.compiler_begin(if_condition, target)\n",
    "        erf_out = relay.erf(cb_5)\n",
    "        ce_4 = relay.annotation.compiler_end(erf_out, target)\n",
    "        func = relay.Function([data, eq1, eq2], ce_4)\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "        return mod\n",
    "\n",
    "    # 获取预期结果并应用类型推断\n",
    "    expected = transform.InferType()(after())\n",
    "    # 测试两种annotate_non_call_ops参数的情况\n",
    "    for annotate_non_call_ops in [True, False]:\n",
    "        # 应用AnnotateTarget转换\n",
    "        result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "        # 验证实际结果与预期结果是否结构相等\n",
    "        tvm.ir.assert_structural_equal(expected, result)\n",
    "\n",
    "\n",
    "def test_while_let():\n",
    "    \"\"\"测试循环和Let绑定在AnnotateTarget转换中的处理\n",
    "\n",
    "    此测试用例验证当计算图包含循环结构和Let绑定时，\n",
    "    AnnotateTarget转换能否正确地处理这些复杂结构并添加编译器注解。\n",
    "    \"\"\"\n",
    "    # 定义目标编译器标识\n",
    "    target = \"test_while_let\"\n",
    "\n",
    "    # 注册less操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"less\", \"target.\" + target)\n",
    "    def less(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册add操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"add\", \"target.\" + target)\n",
    "    def add(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册zeros_like操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"zeros_like\", \"target.\" + target)\n",
    "    def zeros_like(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    def before():\n",
    "        \"\"\"创建包含循环和Let绑定的原始计算图\n",
    "\n",
    "        返回:\n",
    "            包含循环和Let结构的IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        var1 = relay.var(\"var1\", shape=(2,))\n",
    "        var2 = relay.var(\"var2\", shape=(), dtype=\"int32\")\n",
    "        var3 = relay.var(\"var3\", shape=(2,))\n",
    "        # 创建循环条件（var2 < 10）\n",
    "        cond = relay.less(var2, relay.const(10, dtype=\"int32\"))\n",
    "\n",
    "        # 创建循环函数变量\n",
    "        loop = relay.var(\"while_loop\")\n",
    "        # 循环计数器递增\n",
    "        ii = var2 + relay.const(1, dtype=\"int32\")\n",
    "        # 累加计算\n",
    "        ss = var3 + var1\n",
    "        # 递归调用循环函数\n",
    "        true_branch = loop(ii, ss)\n",
    "        # 创建条件分支（循环继续或终止）\n",
    "        ife = relay.If(cond, true_branch, var3)\n",
    "        # 定义循环函数\n",
    "        func_1 = relay.Function([var2, var3], ife)\n",
    "\n",
    "        # 使用Let绑定定义循环函数并启动循环\n",
    "        ret = relay.Let(loop, func_1, loop(relay.const(0, dtype=\"int32\"), relay.zeros_like(var1)))\n",
    "        # 创建主函数\n",
    "        func_2 = relay.Function([var1], ret)\n",
    "        mod = tvm.IRModule.from_expr(func_2)\n",
    "        return mod\n",
    "\n",
    "    def after(annotate_non_call_ops):\n",
    "        \"\"\"生成预期的注解后计算图\n",
    "\n",
    "        参数:\n",
    "            annotate_non_call_ops: 布尔值，表示是否为非调用操作添加注解\n",
    "\n",
    "        返回:\n",
    "            包含正确编译器注解的循环和Let结构IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        var1 = relay.var(\"var1\", shape=(2,))\n",
    "        var2 = relay.var(\"var2\", shape=(), dtype=\"int32\")\n",
    "        var3 = relay.var(\"var3\", shape=(2,))\n",
    "        var4 = relay.const(10, dtype=\"int32\")\n",
    "\n",
    "        # 为循环条件变量添加编译器注解\n",
    "        cb_1 = relay.annotation.compiler_begin(var2, target)\n",
    "        cb_2 = relay.annotation.compiler_begin(var4, target)\n",
    "\n",
    "        # 创建循环条件并添加编译器结束注解\n",
    "        less_condition = relay.less(cb_1, cb_2)\n",
    "        ce_1 = relay.annotation.compiler_end(less_condition, target)\n",
    "\n",
    "        # 创建循环函数变量\n",
    "        loop = relay.var(\"while_loop\")\n",
    "\n",
    "        # 为循环体内的操作添加编译器注解\n",
    "        cb_3 = relay.annotation.compiler_begin(var2, target)\n",
    "        cb_4 = relay.annotation.compiler_begin(relay.const(1, dtype=\"int32\"), target)\n",
    "        add_op_1 = relay.add(cb_3, cb_4)\n",
    "        ce_2 = relay.annotation.compiler_end(add_op_1, target)\n",
    "\n",
    "        # 根据参数决定是否为非调用操作添加注解\n",
    "        cb_5 = relay.annotation.compiler_begin(ce_2, \"default\") if annotate_non_call_ops else ce_2\n",
    "\n",
    "        # 为累加操作添加编译器注解\n",
    "        cb_6 = relay.annotation.compiler_begin(var3, target)\n",
    "        cb_7 = relay.annotation.compiler_begin(var1, target)\n",
    "        add_op_2 = relay.add(cb_6, cb_7)\n",
    "        ce_3 = relay.annotation.compiler_end(add_op_2, target)\n",
    "\n",
    "        # 根据参数决定是否为非调用操作添加注解\n",
    "        cb_8 = relay.annotation.compiler_begin(ce_3, \"default\") if annotate_non_call_ops else ce_3\n",
    "\n",
    "        # 递归调用循环函数（循环体）\n",
    "        true_branch = loop(cb_5, cb_8)  # while loop\n",
    "        # 根据参数决定是否为非调用操作添加注解\n",
    "        ce_4 = (\n",
    "            relay.annotation.compiler_end(true_branch, \"default\")\n",
    "            if annotate_non_call_ops\n",
    "            else true_branch\n",
    "        )\n",
    "        # 创建条件分支\n",
    "        if_condition = relay.If(ce_1, ce_4, var3)\n",
    "        const_1 = relay.const(0, dtype=\"int32\")\n",
    "        # 根据参数决定是否为常量添加注解\n",
    "        cb_9 = (\n",
    "            relay.annotation.compiler_begin(const_1, \"default\")\n",
    "            if annotate_non_call_ops\n",
    "            else const_1\n",
    "        )\n",
    "        # 为zeros_like操作添加编译器注解\n",
    "        cb_10 = relay.annotation.compiler_begin(var1, target)\n",
    "        zeros_like = relay.zeros_like(cb_10)\n",
    "        ce_5 = relay.annotation.compiler_end(zeros_like, target)\n",
    "        # 根据参数决定是否为非调用操作添加注解\n",
    "        cb_11 = relay.annotation.compiler_begin(ce_5, \"default\") if annotate_non_call_ops else ce_5\n",
    "        # 启动循环\n",
    "        while_condition = loop(cb_9, cb_11)\n",
    "        # 根据参数决定是否为非调用操作添加注解\n",
    "        ce_6 = (\n",
    "            relay.annotation.compiler_end(while_condition, \"default\")\n",
    "            if annotate_non_call_ops\n",
    "            else while_condition\n",
    "        )\n",
    "\n",
    "        # 定义循环函数和主函数\n",
    "        func_1 = relay.Function([var2, var3], if_condition)\n",
    "        ret = relay.Let(loop, func_1, ce_6)\n",
    "        func_2 = relay.Function([var1], ret)\n",
    "        mod = tvm.IRModule.from_expr(func_2)\n",
    "        return mod\n",
    "\n",
    "    # 测试两种annotate_non_call_ops参数的情况\n",
    "    for annotate_non_call_ops in [False, True]:\n",
    "        # 应用AnnotateTarget转换\n",
    "        result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "        # 获取预期结果并应用类型推断\n",
    "        expected = transform.InferType()(after(annotate_non_call_ops))\n",
    "        # 验证实际结果与预期结果是否结构相等\n",
    "        tvm.ir.assert_structural_equal(expected, result)\n",
    "\n",
    "\n",
    "def test_if_free_vars():\n",
    "    \"\"\"测试包含自由变量的条件分支在AnnotateTarget转换中的处理\n",
    "\n",
    "    此测试用例验证当条件分支中包含自由变量（如常量生成操作）时，\n",
    "    AnnotateTarget转换能否正确地处理这些结构并添加编译器注解。\n",
    "    \"\"\"\n",
    "    # 定义目标编译器标识\n",
    "    target = \"test_if_free_vars\"\n",
    "\n",
    "    # 注册equal操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"equal\", \"target.\" + target)\n",
    "    def equal(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册sigmoid操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"sigmoid\", \"target.\" + target)\n",
    "    def sigmoid(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    # 注册erf操作在指定目标上的处理函数\n",
    "    @tvm.ir.register_op_attr(\"erf\", \"target.\" + target)\n",
    "    def erf(expr):  # pylint: disable=unused-variable\n",
    "        return True\n",
    "\n",
    "    def before():\n",
    "        \"\"\"创建包含自由变量的条件分支原始计算图\n",
    "\n",
    "        返回:\n",
    "            包含自由变量的条件分支IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        data = relay.var(\"data\", shape=(1, 32))\n",
    "        eq1 = relay.var(\"e1\", shape=[], dtype=\"float32\")\n",
    "        eq2 = relay.var(\"e2\", shape=[], dtype=\"float32\")\n",
    "        # 创建条件表达式\n",
    "        eq = relay.equal(eq1, eq2)\n",
    "\n",
    "        # 条件为真时返回零张量（自由变量示例）\n",
    "        true_branch = relay.zeros(shape=(1, 32), dtype=\"float32\")\n",
    "        # 条件为假时应用sigmoid激活函数\n",
    "        false_branch = relay.sigmoid(data)\n",
    "        # 创建条件分支结构\n",
    "        ife = relay.If(eq, true_branch, false_branch)\n",
    "        # 对条件分支结果应用erf操作\n",
    "        out = relay.erf(ife)\n",
    "\n",
    "        func = relay.Function([data, eq1, eq2], out)\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "\n",
    "        return mod\n",
    "\n",
    "    def after():\n",
    "        \"\"\"生成预期的注解后计算图\n",
    "\n",
    "        返回:\n",
    "            包含正确编译器注解的自由变量条件分支IRModule\n",
    "        \"\"\"\n",
    "        # 创建输入变量\n",
    "        data = relay.var(\"data\", shape=(1, 32))\n",
    "        eq1 = relay.var(\"e1\", shape=[], dtype=\"float32\")\n",
    "        eq2 = relay.var(\"e2\", shape=[], dtype=\"float32\")\n",
    "\n",
    "        # 为条件变量添加编译器注解\n",
    "        cb_1 = relay.annotation.compiler_begin(eq1, target)\n",
    "        cb_2 = relay.annotation.compiler_begin(eq2, target)\n",
    "\n",
    "        # 创建条件表达式并添加编译器结束注解\n",
    "        equality_condition = relay.equal(cb_1, cb_2)\n",
    "        ce_1 = relay.annotation.compiler_end(equality_condition, target)\n",
    "\n",
    "        # 条件为真时返回零张量（自由变量，不需要添加注解）\n",
    "        true_branch = relay.zeros(shape=(1, 32), dtype=\"float32\")\n",
    "\n",
    "        # 为条件为假的分支添加编译器注解\n",
    "        cb_3 = relay.annotation.compiler_begin(data, target)\n",
    "        false_branch = relay.sigmoid(cb_3)\n",
    "        ce_2 = relay.annotation.compiler_end(false_branch, target)\n",
    "\n",
    "        # 创建条件分支结构\n",
    "        if_condition = relay.If(ce_1, true_branch, ce_2)\n",
    "        # 为条件分支结果和最终erf操作添加编译器注解\n",
    "        cb_4 = relay.annotation.compiler_begin(if_condition, target)\n",
    "        erf_out = relay.erf(cb_4)\n",
    "        ce_3 = relay.annotation.compiler_end(erf_out, target)\n",
    "        func = relay.Function([data, eq1, eq2], ce_3)\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "        return mod\n",
    "\n",
    "    # 测试两种annotate_non_call_ops参数的情况\n",
    "    for annotate_non_call_ops in [True, False]:\n",
    "        # 应用AnnotateTarget转换\n",
    "        result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "        # 获取预期结果并应用类型推断\n",
    "        expected = transform.InferType()(after())\n",
    "        # 验证实际结果与预期结果是否结构相等\n",
    "        tvm.ir.assert_structural_equal(expected, result)\n",
    "\n",
    "\n",
    "def test_free_vars_zeros():\n",
    "    \"\"\"测试自由变量在AnnotateTarget转换中的独立处理\n",
    "\n",
    "    此测试用例验证自由变量（如零张量生成操作）在单独存在时，\n",
    "    AnnotateTarget转换能否正确地处理它们。\n",
    "    \"\"\"\n",
    "    # 定义目标编译器标识\n",
    "    target = \"test_free_vars_zeros\"\n",
    "\n",
    "    def before():\n",
    "        \"\"\"创建包含自由变量的原始计算图\n",
    "\n",
    "        返回:\n",
    "            仅包含零张量生成操作的IRModule\n",
    "        \"\"\"\n",
    "        # 创建一个不接受参数，仅返回零张量的函数\n",
    "        func = relay.Function([], relay.zeros(shape=(0), dtype=\"float32\"))\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "        return mod\n",
    "\n",
    "    def after():\n",
    "        \"\"\"生成预期的注解后计算图\n",
    "\n",
    "        返回:\n",
    "            预期的不包含编译器注解的零张量IRModule\n",
    "        \"\"\"\n",
    "        # 自由变量不应被添加编译器注解，因此预期结果与原始结果相同\n",
    "        func = relay.Function([], relay.zeros(shape=(0), dtype=\"float32\"))\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "        return mod\n",
    "\n",
    "    # 应用AnnotateTarget转换\n",
    "    result = transform.AnnotateTarget(target)(before())\n",
    "    # 获取预期结果并应用类型推断\n",
    "    expected = transform.InferType()(after())\n",
    "    # 验证实际结果与预期结果是否结构相等\n",
    "    tvm.ir.assert_structural_equal(expected, result)\n",
    "\n",
    "\n",
    "def test_empty_tuple():\n",
    "    \"\"\"测试空元组在AnnotateTarget转换中的处理\n",
    "\n",
    "    此测试用例验证空元组是否能像无参数调用一样正确地被AnnotateTarget转换处理。\n",
    "    \"\"\"\n",
    "    # 定义目标编译器标识\n",
    "    target = \"test_empty_tuple\"\n",
    "\n",
    "    def before():\n",
    "        \"\"\"创建包含空元组的原始计算图\n",
    "\n",
    "        返回:\n",
    "            仅包含空元组的IRModule\n",
    "        \"\"\"\n",
    "        # 创建一个不接受参数，仅返回空元组的函数\n",
    "        func = relay.Function([], relay.Tuple([]))\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "        return mod\n",
    "\n",
    "    def after():\n",
    "        \"\"\"生成预期的注解后计算图\n",
    "\n",
    "        返回:\n",
    "            预期的不包含编译器注解的空元组IRModule\n",
    "        \"\"\"\n",
    "        # 空元组不应被添加编译器注解，因此预期结果与原始结果相同\n",
    "        func = relay.Function([], relay.Tuple([]))\n",
    "        mod = tvm.IRModule.from_expr(func)\n",
    "        return mod\n",
    "\n",
    "    # 测试两种annotate_non_call_ops参数的情况\n",
    "    for annotate_non_call_ops in [True, False]:\n",
    "        # 应用AnnotateTarget转换\n",
    "        result = transform.AnnotateTarget(target, annotate_non_call_ops)(before())\n",
    "        # 获取预期结果并应用类型推断\n",
    "        expected = transform.InferType()(after())\n",
    "        # 验证实际结果与预期结果是否结构相等\n",
    "        tvm.ir.assert_structural_equal(expected, result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db07a9b5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
